import os
from dotenv import load_dotenv
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, declarative_base

# 加载环境变量
load_dotenv(dotenv_path="application.env")


# 定义一个函数来创建数据库引擎
def create_db_engine(prefix=""):
    """创建数据库引擎，prefix用于区分不同环境的变量前缀（如LOCAL_）"""
    DB_USER = "postgres"
    DB_PASSWORD = "Yilun123456"
    DB_HOST = "180.166.193.6"
    # DB_HOST = "192.168.2.107"
    DB_PORT = 55432
    DB_NAME = "lab"
    # 获取环境变量，不进行encode()转换
    db_user = os.environ.get(f"{prefix}DB_USER", "")
    db_password = os.environ.get(f"{prefix}DB_PASSWORD", "")
    db_host = os.environ.get(f"{prefix}DB_HOST", "")
    # 确保端口号是整数
    db_port = int(os.environ.get(f"{prefix}DB_PORT", 55432))
    db_name = os.environ.get(f"{prefix}DB_NAME", "")

    # 构建数据库URL
    database_url = f"postgresql+psycopg2://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

    # 创建并返回数据库引擎
    return create_engine(
        database_url,
        echo=True,
        pool_size=5,
        max_overflow=10,
        pool_timeout=30,
    )


# 创建生产环境数据库连接
engine = create_db_engine()

# 创建会话工厂
SessionLocal = sessionmaker(
    bind=engine,
    autocommit=False,
    autoflush=False,
)

# 创建基类
Base = declarative_base()


# 获取数据库会话（用于依赖注入）
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


# 创建本地开发环境数据库连接
local_engine = create_db_engine(prefix="LOCAL_")

# 创建本地会话工厂
LocalSessionLocal = sessionmaker(
    bind=local_engine,
    autocommit=False,
    autoflush=False,
)


# 获取本地数据库会话
def get_local_db():
    db = LocalSessionLocal()
    try:
        yield db
    finally:
        db.close()
